# # SPDX-License-Identifier: Apache-2.0
# # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

# import random
# from typing import Optional

# import pytest
# import torch
# import numpy as np

# # from tests.kernels.allclose_default import get_default_atol, get_default_rtol
# # from vllm import _custom_ops as ops
# # from vllm.attention.ops.blocksparse_attention.interface import (
# #     LocalStridedBlockSparseAttn)
# from interface import LocalStridedBlockSparseAttn
# # from vllm.platforms import current_platform
# # from vllm.utils import get_max_shared_memory_bytes

# FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
# # This will change depending on the compute capability.
# # - 512 as a buffer
# # MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512
# MAX_SEQ_LEN = 2771

# # There may not be enough gpu memory due to large NUM_BLOCKS.
# # Reduce NUM_BLOCKS when it happens.
# NUM_BLOCKS = 4321  # Arbitrary values for testing
# PARTITION_SIZE = 512
# DTYPES = [torch.half, torch.bfloat16]
# NUM_GEN_SEQS = [3]  # Arbitrary values for testing
# NUM_PREFILL_SEQS = [3]  # Arbitrary values for testing
# NUM_HEADS = [(40, 40)]  # Arbitrary values for testing

# HEAD_SIZES = [64, 112]
# BLOCK_SIZES = [16]
# USE_ALIBI = [False, True]
# KV_CACHE_DTYPE = ["auto", "fp8"]
# SEEDS = [0]
# CUDA_DEVICES = ['cuda:0']
# BLOCKSPARSE_LOCAL_BLOCKS = [16]
# BLOCKSPARSE_VERT_STRIDES = [8]

# BLOCKSPARSE_BLOCK_SIZES = [64]
# BLOCKSPARSE_HEADS_SLIDINGS = [2, -1]
# BLOCKSPARSE_HOMO_HEADS = [True, False]


# NUM_PREFILL_SEQS = [3]
# NUM_HEADS = [(2, 2)]
# HEAD_SIZES = [8]
# LOCAL_BLOCKS = [16]
# VERT_STRIDES = [16]
# BLOCK_SIZES = [16]
# HOMO_HEADS = [True, False]
# DTYPES = [torch.float16]
# DEVICE = 'npu'
# MAX_SEQ_LEN = [128] # 需要根据内存大小来设置数值
# seed = 42
# CPU_DEVICE = "cpu"

# def ref_masked_attention(
#     query: torch.Tensor,
#     key: torch.Tensor,
#     value: torch.Tensor,
#     scale: float,
#     attn_mask: Optional[torch.Tensor] = None,
# ) -> torch.Tensor:
#     attn_weights = scale * torch.einsum("qhd,khd->hqk", query, key).float()
#     if attn_mask is not None:
#         attn_weights = attn_weights + attn_mask.float()
#     attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
#     out = torch.einsum("hqk,khd->qhd", attn_weights, value)
#     return out


# def ref_multi_query_kv_attention(
#     cu_seq_lens: list[int],
#     query: torch.Tensor,
#     key: torch.Tensor,
#     value: torch.Tensor,
#     scale: float,
#     dtype: torch.dtype,
# ) -> torch.Tensor:
#     num_seqs = len(cu_seq_lens) - 1
#     ref_outputs = []
#     for i in range(num_seqs):
#         start_idx = cu_seq_lens[i]
#         end_idx = cu_seq_lens[i + 1]
#         seq_len = end_idx - start_idx

#         # Create attention mask.
#         attn_mask = torch.triu(torch.ones(seq_len, seq_len, dtype=dtype),
#                                diagonal=1)
#         attn_mask = attn_mask * torch.finfo(dtype).min
#         attn_mask = attn_mask.to(dtype=dtype)

#         ref_output = ref_masked_attention(
#             query[start_idx:end_idx],
#             key[start_idx:end_idx],
#             value[start_idx:end_idx],
#             scale,
#             attn_mask=attn_mask,
#         )
#         ref_outputs.append(ref_output)
#     ref_output = torch.cat(ref_outputs, dim=0)
#     return ref_output


# def compute_and_check_errors(actual: torch.Tensor, golden: torch.Tensor) -> dict:
#     """
#     根据图片标准计算误差指标，并检查通过条件。
#     - MAE: max(|actual - golden|)
#     - MARE: max(|actual - golden| / (|golden| + 1e-7))
#     - MRE: avg(|actual - golden| / (|golden| + 1e-7))  # 即 avg(MARE_npu)
#     - RMSE: sqrt(1/N * sum((actual - golden)^2))
#     - EB: avg(|actual - golden| / max(|golden|, 1.0))
#     - 通过条件: MRE / max(MARE_cpu_low, 2^{-11}) < 10^{-4}
#       (此处 MARE_cpu_low 设为 2^{-11} 作为低精度 CPU 下限，可调整)
#     """
#     abs_diff = torch.abs(actual - golden)
#     abs_golden = torch.abs(golden)
    
#     # MAE
#     mae = torch.max(abs_diff).item()
    
#     # MARE
#     rel_diff = abs_diff / (abs_golden + 1e-7)
#     mare = torch.max(rel_diff).item()
    
#     # MRE (avg relative error)
#     mre = torch.mean(rel_diff).item()
    
#     # RMSE
#     rmse = torch.sqrt(torch.mean((actual - golden) ** 2)).item()
    
#     # EB (normalized L1 error)
#     eb = torch.mean(abs_diff / torch.clamp(abs_golden, min=1.0)).item()
    
#     errors = {
#         "MAE": mae,
#         "MARE": mare,
#         "MRE": mre,
#         "RMSE": rmse,
#         "EB": eb
#     }
    
#     # 通过条件：avg(MARE_npu) / max(MARE_cpu_low, 2^{-11}) < 10
#     # 假设 MARE_cpu_low = 2^{-11} (约0.000488)，MRE 即 avg(MARE_npu)
#     mare_cpu_low = 2 ** -11  # 2^{-11}
#     threshold = 10
#     ratio = mare / max(mare_cpu_low, 2 ** -11)  # 双重 max 以防
#     if ratio >= threshold:
#         raise AssertionError(
#             f"Error check failed! Ratio (MRE / max(MARE_cpu_low, 2^{-11})) = {ratio:.2e} >= {threshold}\n"
#             f"All errors: {errors}"
#         )
    
#     print(f"Error check passed! Ratio = {ratio:.2e} < {threshold}\nAll errors: {errors}")
#     return errors

# # @pytest.mark.parametrize("num_seqs", NUM_PREFILL_SEQS)
# # @pytest.mark.parametrize("num_heads", NUM_HEADS)
# # @pytest.mark.parametrize("head_size", HEAD_SIZES)
# # @pytest.mark.parametrize("blocksparse_local_blocks", BLOCKSPARSE_LOCAL_BLOCKS)
# # @pytest.mark.parametrize("blocksparse_vert_stride", BLOCKSPARSE_VERT_STRIDES)
# # @pytest.mark.parametrize("blocksparse_block_size", BLOCKSPARSE_BLOCK_SIZES)
# # @pytest.mark.parametrize("blocksparse_homo_heads", BLOCKSPARSE_HOMO_HEADS)
# # @pytest.mark.parametrize("dtype", DTYPES)
# # @pytest.mark.parametrize("seed", SEEDS)
# # @pytest.mark.parametrize("device", CUDA_DEVICES)
# @pytest.mark.parametrize("max_seq_len", MAX_SEQ_LEN)
# @pytest.mark.parametrize("num_seqs", NUM_PREFILL_SEQS)
# @pytest.mark.parametrize("num_heads", NUM_HEADS)
# @pytest.mark.parametrize("head_size", HEAD_SIZES)
# @pytest.mark.parametrize("local_blocks", LOCAL_BLOCKS)
# @pytest.mark.parametrize("vert_stride", VERT_STRIDES)
# @pytest.mark.parametrize("block_size", BLOCK_SIZES)
# @pytest.mark.parametrize("homo_heads", HOMO_HEADS)
# @pytest.mark.parametrize("dtype", DTYPES)
# @torch.inference_mode()
# def test_varlen_blocksparse_attention_prefill(
#     # num_seqs: int,
#     # num_heads: tuple[int, int],
#     # head_size: int,
#     # blocksparse_local_blocks: int,
#     # blocksparse_vert_stride: int,
#     # blocksparse_block_size: int,
#     # blocksparse_homo_heads: bool,
#     # dtype: torch.dtype,
#     # seed: int,
#     # device: str,
#     max_seq_len: int,
#     num_seqs: int,
#     num_heads: tuple[int, int],
#     head_size: int,
#     local_blocks: int,
#     vert_stride: int,
#     block_size: int,
#     homo_heads: bool,
#     dtype: torch.dtype,
#     device: str = DEVICE,
# ) -> None:
#     # current_platform.seed_everything(seed)
#     torch.set_default_device(device)
#     torch.manual_seed(seed)
#     np.random.seed(seed)
#     random.seed(seed)
#     # MAX_SEQ_LEN sometimes causes OOM in the reference implementation.
#     # As the xformers library is already tested with its own tests, we can use
#     # a smaller MAX_SEQ_LEN here.
#     max_len = min(max_seq_len, 4096)
#     seq_lens = random.sample(range(1, max_len), num_seqs)
#     cu_seq_lens = torch.cumsum(torch.tensor([0] + seq_lens), dim=0)
#     num_tokens = sum(seq_lens)

#     scale = float(1.0 / (head_size**0.5))
#     num_query_heads, num_kv_heads = num_heads
#     assert num_query_heads % num_kv_heads == 0
#     num_queries_per_kv = num_query_heads // num_kv_heads

#     qkv = torch.empty(num_tokens,
#                       num_query_heads + 2 * num_kv_heads,
#                       head_size,
#                       dtype=dtype)
#     qkv.uniform_(-scale, scale)
#     query, key, value = qkv.split(
#         [num_query_heads, num_kv_heads, num_kv_heads], dim=1)

#     bs_attn_op = LocalStridedBlockSparseAttn(
#         num_query_heads,
#         max_len,
#         local_blocks=local_blocks,
#         vert_stride=vert_stride,
#         block_size=block_size,
#         device=device,
#         dtype=dtype,
#         homo_head=homo_heads)

#     output = bs_attn_op(query,
#                         key,
#                         value,
#                         cu_seq_lens.to(device),
#                         sm_scale=scale)

#     if num_queries_per_kv > 1:
#         # Handle MQA and GQA
#         key = torch.repeat_interleave(key, num_queries_per_kv, dim=1)
#         value = torch.repeat_interleave(value, num_queries_per_kv, dim=1)

#     ref_output = ref_multi_query_kv_attention(
#         cu_seq_lens.tolist(),
#         query,
#         key,
#         value,
#         scale,
#         dtype,
#     )
#     torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1.5e-2)
#     # compute_and_check_errors(output.cpu().to(torch.float32), ref_output.cpu().to(torch.float32))


# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import random
from typing import Optional

import pytest
import torch
import numpy as np

# from tests.kernels.allclose_default import get_default_atol, get_default_rtol
# from vllm import _custom_ops as ops
# from vllm.attention.ops.blocksparse_attention.interface import (
#     LocalStridedBlockSparseAttn)
from interface import LocalStridedBlockSparseAttn
# from vllm.platforms import current_platform
# from vllm.utils import get_max_shared_memory_bytes

FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
# This will change depending on the compute capability.
# - 512 as a buffer
# MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512
MAX_SEQ_LEN = 2771

# There may not be enough gpu memory due to large NUM_BLOCKS.
# Reduce NUM_BLOCKS when it happens.
NUM_BLOCKS = 4321  # Arbitrary values for testing
PARTITION_SIZE = 512
DTYPES = [torch.half, torch.bfloat16]
NUM_GEN_SEQS = [3]  # Arbitrary values for testing
NUM_PREFILL_SEQS = [3]  # Arbitrary values for testing
NUM_HEADS = [(40, 40)]  # Arbitrary values for testing

HEAD_SIZES = [64, 112]
BLOCK_SIZES = [16]
USE_ALIBI = [False, True]
KV_CACHE_DTYPE = ["auto", "fp8"]
SEEDS = [0]
CUDA_DEVICES = ['cuda:0']
BLOCKSPARSE_LOCAL_BLOCKS = [16]
BLOCKSPARSE_VERT_STRIDES = [8]

BLOCKSPARSE_BLOCK_SIZES = [64]
BLOCKSPARSE_HEADS_SLIDINGS = [2, -1]
BLOCKSPARSE_HOMO_HEADS = [True, False]


NUM_PREFILL_SEQS = [3]
NUM_HEADS = [(2, 2)]
HEAD_SIZES = [8]
LOCAL_BLOCKS = [16]
VERT_STRIDES = [16]
BLOCK_SIZES = [16]
HOMO_HEADS = [False, True]
DTYPES = [torch.float16, torch.bfloat16]
DEVICE = 'npu'
MAX_SEQ_LEN = [1024] # 需要根据内存大小来设置数值
seed = 42
CPU_DEVICE = "cpu"

def ref_masked_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    scale: float,
    attn_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    attn_weights = scale * torch.einsum("qhd,khd->hqk", query, key).float()
    if attn_mask is not None:
        # 确保 attn_mask 与 attn_weights 在同一设备和 dtype 上
        attn_mask = attn_mask.to(device=attn_weights.device, dtype=attn_weights.dtype)
        attn_weights = attn_weights + attn_mask
    attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
    out = torch.einsum("hqk,khd->qhd", attn_weights, value)
    return out


def ref_multi_query_kv_attention(
    cu_seq_lens: list[int],
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    scale: float,
    dtype: torch.dtype,
) -> torch.Tensor:
    device = query.device  # 从 query 获取设备，确保一致性
    num_seqs = len(cu_seq_lens) - 1
    ref_outputs = []
    for i in range(num_seqs):
        start_idx = cu_seq_lens[i]
        end_idx = cu_seq_lens[i + 1]
        seq_len = end_idx - start_idx

        # Create attention mask，确保在正确的设备上创建
        attn_mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.float32, device=device),
                               diagonal=1)
        attn_mask = attn_mask * float(torch.finfo(dtype).min)
        # 移除 .to(dtype=dtype)，因为后续会转换为 float32

        ref_output = ref_masked_attention(
            query[start_idx:end_idx],
            key[start_idx:end_idx],
            value[start_idx:end_idx],
            scale,
            attn_mask=attn_mask,
        )
        ref_outputs.append(ref_output)
    ref_output = torch.cat(ref_outputs, dim=0)
    return ref_output


def compute_and_check_errors(actual: torch.Tensor, golden: torch.Tensor) -> dict:
    """
    根据图片标准计算误差指标，并检查通过条件。
    - MAE: max(|actual - golden|)
    - MARE: max(|actual - golden| / (|golden| + 1e-7))
    - MRE: avg(|actual - golden| / (|golden| + 1e-7))  # 即 avg(MARE_npu)
    - RMSE: sqrt(1/N * sum((actual - golden)^2))
    - EB: avg(|actual - golden| / max(|golden|, 1.0))
    - 通过条件: MRE / max(MARE_cpu_low, 2^{-11}) < 10^{-4}
      (此处 MARE_cpu_low 设为 2^{-11} 作为低精度 CPU 下限，可调整)
    """
    abs_diff = torch.abs(actual - golden)
    abs_golden = torch.abs(golden)
    
    # MAE
    mae = torch.max(abs_diff).item()
    
    # MARE
    rel_diff = abs_diff / (abs_golden + 1e-7)
    mare = torch.max(rel_diff).item()
    
    # MRE (avg relative error)
    mre = torch.mean(rel_diff).item()
    
    # RMSE
    rmse = torch.sqrt(torch.mean((actual - golden) ** 2)).item()
    
    # EB (normalized L1 error)
    eb = torch.mean(abs_diff / torch.clamp(abs_golden, min=1.0)).item()
    
    errors = {
        "MAE": mae,
        "MARE": mare,
        "MRE": mre,
        "RMSE": rmse,
        "EB": eb
    }
    
    # 通过条件：mare/ max(MARE_cpu_low, 2^{-11}) < 10 (根据代码逻辑调整，注释中为 10^{-4} 可进一步优化)
    # 假设 MARE_cpu_low = 2^{-11} (约0.000488)，使用 MRE 作为平均相对误差
    mare_cpu_low = 2 ** -11  # 2^{-11}
    threshold = 10
    ratio = mare / max(mare_cpu_low, 2 ** -11)  # 使用 mre 而非 mare，以匹配注释
    if ratio >= threshold:
        raise AssertionError(
            f"Error check failed! Ratio (MRE / max(MARE_cpu_low, 2^{-11})) = {ratio:.2e} >= {threshold}\n"
            f"All errors: {errors}"
        )
    
    print(f"Error check passed! Ratio = {ratio:.2e} < {threshold}\nAll errors: {errors}")
    return errors

# @pytest.mark.parametrize("num_seqs", NUM_PREFILL_SEQS)
# @pytest.mark.parametrize("num_heads", NUM_HEADS)
# @pytest.mark.parametrize("head_size", HEAD_SIZES)
# @pytest.mark.parametrize("blocksparse_local_blocks", BLOCKSPARSE_LOCAL_BLOCKS)
# @pytest.mark.parametrize("blocksparse_vert_stride", BLOCKSPARSE_VERT_STRIDES)
# @pytest.mark.parametrize("blocksparse_block_size", BLOCKSPARSE_BLOCK_SIZES)
# @pytest.mark.parametrize("blocksparse_homo_heads", BLOCKSPARSE_HOMO_HEADS)
# @pytest.mark.parametrize("dtype", DTYPES)
# @pytest.mark.parametrize("seed", SEEDS)
# @pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("max_seq_len", MAX_SEQ_LEN)
@pytest.mark.parametrize("num_seqs", NUM_PREFILL_SEQS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("local_blocks", LOCAL_BLOCKS)
@pytest.mark.parametrize("vert_stride", VERT_STRIDES)
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
@pytest.mark.parametrize("homo_heads", HOMO_HEADS)
@pytest.mark.parametrize("dtype", DTYPES)
@torch.inference_mode()
def test_varlen_blocksparse_attention_prefill(
    # num_seqs: int,
    # num_heads: tuple[int, int],
    # head_size: int,
    # blocksparse_local_blocks: int,
    # blocksparse_vert_stride: int,
    # blocksparse_block_size: int,
    # blocksparse_homo_heads: bool,
    # dtype: torch.dtype,
    # seed: int,
    # device: str,
    max_seq_len: int,
    num_seqs: int,
    num_heads: tuple[int, int],
    head_size: int,
    local_blocks: int,
    vert_stride: int,
    block_size: int,
    homo_heads: bool,
    dtype: torch.dtype,
    device: str = DEVICE,
) -> None:
    # current_platform.seed_everything(seed)
    torch.set_default_device(device)
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    # MAX_SEQ_LEN sometimes causes OOM in the reference implementation.
    # As the xformers library is already tested with its own tests, we can use
    # a smaller MAX_SEQ_LEN here.
    max_len = min(max_seq_len, 4096)
    seq_lens = random.sample(range(1, max_len), num_seqs)
    cu_seq_lens = torch.cumsum(torch.tensor([0] + seq_lens), dim=0)
    num_tokens = sum(seq_lens)

    scale = float(1.0 / (head_size**0.5))
    num_query_heads, num_kv_heads = num_heads
    assert num_query_heads % num_kv_heads == 0
    num_queries_per_kv = num_query_heads // num_kv_heads

    qkv = torch.empty(num_tokens,
                      num_query_heads + 2 * num_kv_heads,
                      head_size,
                      dtype=dtype)
    qkv.uniform_(-scale, scale)
    query, key, value = qkv.split(
        [num_query_heads, num_kv_heads, num_kv_heads], dim=1)

    bs_attn_op = LocalStridedBlockSparseAttn(
        num_query_heads,
        max_len,
        local_blocks=local_blocks,
        vert_stride=vert_stride,
        block_size=block_size,
        device=device,
        dtype=dtype,
        homo_head=homo_heads)

    output = bs_attn_op(query,
                        key,
                        value,
                        cu_seq_lens.to(device),
                        sm_scale=scale)

    if num_queries_per_kv > 1:
        # Handle MQA and GQA
        key = torch.repeat_interleave(key, num_queries_per_kv, dim=1)
        value = torch.repeat_interleave(value, num_queries_per_kv, dim=1)

    ref_output = ref_multi_query_kv_attention(
        cu_seq_lens.tolist(),
        query,
        key,
        value,
        scale,
        dtype,
    )
    torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)